import cv2
import os
from tqdm import tqdm
import numpy as np


def getAllImage(folderPath, imageList, oldPath, newPath):
    extend_name = ["jpg", "jpeg", "png", "bmp"]
    for item in os.listdir(folderPath):
        if os.path.isdir(os.path.join(folderPath, item)):
            subFolderPath = os.path.join(folderPath, item)
            newFolderPath = subFolderPath.replace(oldPath, newPath)
            if not os.path.exists(newFolderPath):
                os.mkdir(newFolderPath)
            getAllImage(subFolderPath, imageList, oldPath, newPath)
        else:

            filePath = os.path.join(folderPath, item)
            if os.path.isfile(filePath):
                if item.split('.')[-1] in extend_name:
                    imageList.append(filePath)
    return imageList


def create_mask(imgpath, threshold):
    image = cv2.imread(imgpath, cv2.IMREAD_GRAYSCALE)
    _, mask = cv2.threshold(image, threshold, 255, cv2.THRESH_BINARY)
    return mask


def check(image, mask, bright):
    height = mask.shape[0]
    weight = mask.shape[1]
    for row in tqdm(range(height)):  # 遍历高

        for col in range(weight):  # 遍历宽
            if (mask[row, col] > 200):
                if bright == 3:
                    image[row, col] = [193, 207, 221]
                if bright == 2:
                    image[row, col] = [142, 157, 163]
                if bright == 1:
                    image[row, col] = [122, 133, 135]
                if bright == 0:
                    image[row, col] = [71, 80, 86]
    return image


def checkBright(image):
    # averageColors = np.average(image, axis=(0, 1))
    # averageColor = averageColors[0] + averageColors[1] + averageColors[2]
    color_thd = 50
    red = []
    green = []
    blue = []
    for row in tqdm(range(image.shape[0])):
        for col in range(image.shape[1]):
            r, g, b = image[row, col]
            if r <= color_thd and g <= color_thd and b <= color_thd:
                continue
            if r <= color_thd and g <= color_thd and b <= color_thd:
                continue
            red.append(r)
            green.append(g)
            blue.append(b)
    aveRed = np.average(np.array(red))
    aveGreen = np.average(np.array(green))
    aveBlue = np.average(np.array(blue))
    averageColor = aveRed + aveGreen + aveBlue
    averageColor /= 3
    print(averageColor)
    if averageColor > 150:
        return 3
    if 150 > averageColor > 120:
        return 2
    if 90 < averageColor < 120:
        return 1
    if averageColor < 90:
        return 0


if __name__ == '__main__':
    path = "E:\dataset1\\2020"
    newPath = "E:\dataset2\\2020"
    imageList = ["1/Image_20210928110616811.bmp"]
    rootpath = "data"
    masksavepath = "mask"
    savepath = "result"
    imgfiles = os.listdir(rootpath)
    #imageList = getAllImage(path, imageList, path, newPath)

    brightthreshold = 200
    darkthreshold = 140
    darkthresholdMore=100
    for i in range(0, len(imageList)):
        if "草地贪夜蛾雄虫" in imageList[i]:
            imageList[i]=imageList[i].replace("草地贪夜蛾雄虫", "")
        if "草地贪夜蛾" in imageList[i]:
            imageList[i]=imageList[i].replace("草地贪夜蛾", "")
        #imageList[i]=imageList[i].replace("1/", "")
        imagefile = imageList[i].split("/")[-1]

        print(imagefile)
        if imagefile.endswith("jpg") or imagefile.endswith("JPG") or imagefile.endswith("bmp"):
            image = cv2.imread(imageList[i])
            maskpath = os.path.join(masksavepath, "mask_" + imagefile)

            bright = checkBright(image)
            if bright == 3:
                mask = create_mask(imageList[i], brightthreshold)
                # cv2.imwrite(maskpath, mask)

                xiufuImage = check(image, mask, bright)
                newImagePath = imageList[i].replace(path, newPath)
                cv2.imwrite(newImagePath, xiufuImage)

            if bright == 2:
                mask = create_mask(imageList[i], brightthreshold)
                # cv2.imwrite(maskpath, mask)

                xiufuImage = check(image, mask, bright)
                newImagePath = imageList[i].replace(path, newPath)
                cv2.imwrite(newImagePath, xiufuImage)
            if bright == 1:
                mask = create_mask(imageList[i], darkthreshold)
                # cv2.imwrite(maskpath, mask)

                xiufuImage = check(image, mask, bright)
                newImagePath = imageList[i].replace(path, newPath)
                cv2.imwrite(newImagePath, xiufuImage)

            if bright == 0:
                mask = create_mask(imageList[i], darkthresholdMore)
                cv2.imwrite(maskpath, mask)

                xiufuImage = check(image, mask, bright)
                newImagePath = imageList[i].replace(path, newPath)
                #cv2.imwrite(newImagePath, xiufuImage)
                cv2.imwrite(os.path.join(savepath, "xiufu_" + imagefile), xiufuImage)
